package com.demo.spring.rag;

import java.util.List;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.loader.ClassPathDocumentLoader;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.apache.pdfbox.ApachePdfBoxDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;

@Configuration
public class CommonConfig {

	@Autowired
	private OpenAiChatModel model;
	@Autowired
	private ChatMemoryStore store;
	@Autowired
	private EmbeddingModel embeddingModel;
	
	// 构建会话记忆体
	@Bean
	public ChatMemoryProvider chatMemoryProvider() {
		ChatMemoryProvider chatMemoryProvider = new ChatMemoryProvider() {

			public ChatMemory get(Object memoryId) {
				MessageWindowChatMemory memory = MessageWindowChatMemory.builder()
						.id(memoryId)
						.chatMemoryStore(store)
						.maxMessages(20)
						.build();
				return memory;
			}
			
		};
		return chatMemoryProvider;
	}
	
	// 构建向量数据库
	@Bean
	@Primary
	public EmbeddingStore store() {
		// 1、加载文件进内存
		List<Document> documents = ClassPathDocumentLoader.loadDocuments("document");
		// 2、构建向量数据库操作对象
		InMemoryEmbeddingStore store = new InMemoryEmbeddingStore();
		// 3、构建文档分割器
		DocumentSplitter ds = DocumentSplitters.recursive(300, 100);
		// 4、构建一个EmbeddingStoreIngestor对象，完成文本数据切割，向量化，存储
		EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
				.embeddingStore(store)
				.documentSplitter(ds)
				.embeddingModel(embeddingModel)
				.build();
		ingestor.ingest(documents);
		return store;
	}
	
	// 构建向量数据库检索对象
	@Bean
	public ContentRetriever contentRetriever(EmbeddingStore store) {
		return EmbeddingStoreContentRetriever.builder()
				.embeddingStore(store)
				.minScore(0.5)
				.maxResults(3)
				.embeddingModel(embeddingModel)
				.build();
	}
	
}
